﻿#pragma once
#include <iostream>
#include <algorithm>
#include <functional>
#include "types.hpp"
class ModbusSlave {
public:
    using ExceptionCallback = std::function<void(uint8_t, uint8_t, uint8_t)>;

    ModbusSlave(uint8_t slave_id, ModbusData& data)
        : slave_id_(slave_id), data_(data) {
    }

    std::vector<uint8_t> processRequest(const std::vector<uint8_t>& request, ProtocolType protocol) {
        if (request.empty()) {
            std::cerr << "收到空请求" << std::endl;
            return std::vector<uint8_t>();
        }
        // 对于Modbus TCP，提取事务ID
        uint16_t transaction_id = 0;
        if (protocol == ProtocolType::MODBUS_TCP && request.size() >= 2) {
            transaction_id = (request[0] << 8) | request[1];
        }

        try {

            // 移除协议头，获取PDU
            std::vector<uint8_t> pdu = removeProtocolHeader(request, protocol);
            if (pdu.empty()) {
                std::cerr << "协议头解析失败或CRC校验失败" << std::endl;
                return std::vector<uint8_t>();
            }

            // 检查从站ID
            if (pdu[0] != slave_id_) {
                // 不是给这个从站的请求，静默忽略
                return std::vector<uint8_t>();
            }

            if (pdu.size() < 2) {
                std::cerr << "PDU长度不足" << std::endl;
                return createExceptionResponse(0, 0x04, protocol, transaction_id); // 从站设备故障
            }

            uint8_t function_code = pdu[1];
            std::vector<uint8_t> response;

            switch (function_code) {
            case 0x01: // 读线圈
                response = handleReadCoils(pdu, protocol, transaction_id);
                break;
            case 0x02: // 读离散输入
                response = handleReadDiscreteInputs(pdu, protocol, transaction_id);
                break;
            case 0x03: // 读保持寄存器
                response = handleReadHoldingRegisters(pdu, protocol, transaction_id);
                break;
            case 0x04: // 读输入寄存器
                response = handleReadInputRegisters(pdu, protocol, transaction_id);
                break;
            case 0x05: // 写单个线圈
                response = handleWriteSingleCoil(pdu, protocol, transaction_id);
                break;
            case 0x06: // 写单个寄存器
                response = handleWriteSingleRegister(pdu, protocol, transaction_id);
                break;
            case 0x0F: // 写多个线圈
                response = handleWriteMultipleCoils(pdu, protocol, transaction_id);
                break;
            case 0x10: // 写多个寄存器
                response = handleWriteMultipleRegisters(pdu, protocol, transaction_id);
                break;
            default:
                std::cerr << "不支持的功能码: 0x" << std::hex << static_cast<int>(function_code) << std::endl;
                response = createExceptionResponse(function_code, 0x01, protocol, transaction_id); // 非法功能
                break;
            }

            return response;

        }
        catch (const std::exception& e) {
            std::cerr << "处理请求时发生异常: " << e.what() << std::endl;
            return createExceptionResponse(0, 0x04, protocol, transaction_id); // 从站设备故障
        }
    }

private:
    std::vector<uint8_t> handleReadCoils(const std::vector<uint8_t>& request, ProtocolType protocol, uint16_t transaction_id) {
        if (request.size() < 6) {
            return createExceptionResponse(0x01, 0x03, protocol, transaction_id); // 非法数据值
        }

        uint16_t start_address = (request[2] << 8) | request[3];
        uint16_t quantity = (request[4] << 8) | request[5];

        if (quantity < 1 || quantity > 2000) {
            return createExceptionResponse(0x01, 0x03, protocol, transaction_id);
        }

        if (!validateAddressRange(start_address, quantity, data_.coils.size())) {
            return createExceptionResponse(0x01, 0x02, protocol, transaction_id); // 非法数据地址
        }

        // 构建响应
        std::vector<uint8_t> response_pdu;
        response_pdu.push_back(slave_id_);
        response_pdu.push_back(0x01);

        uint8_t byte_count = static_cast<uint8_t>((quantity + 7) / 8);
        response_pdu.push_back(byte_count);

        // 打包线圈状态到字节
        for (uint16_t byte_index = 0; byte_index < byte_count; ++byte_index) {
            uint8_t coil_byte = 0;
            for (uint8_t bit_index = 0; bit_index < 8; ++bit_index) {
                uint16_t coil_index = start_address + byte_index * 8 + bit_index;
                if (coil_index < start_address + quantity && coil_index < data_.coils.size()) {
                    if (data_.coils[coil_index]) {
                        coil_byte |= (1 << bit_index);
                    }
                }
            }
            response_pdu.push_back(coil_byte);
        }

        return addProtocolHeader(response_pdu, protocol, transaction_id);
    }

    std::vector<uint8_t> handleReadDiscreteInputs(const std::vector<uint8_t>& request, ProtocolType protocol, uint16_t transaction_id) {
        if (request.size() < 6) {
            return createExceptionResponse(0x02, 0x03, protocol, transaction_id);
        }

        uint16_t start_address = (request[2] << 8) | request[3];
        uint16_t quantity = (request[4] << 8) | request[5];

        if (quantity < 1 || quantity > 2000) {
            return createExceptionResponse(0x02, 0x03, protocol, transaction_id);
        }

        if (!validateAddressRange(start_address, quantity, data_.discrete_inputs.size())) {
            return createExceptionResponse(0x02, 0x02, protocol, transaction_id);
        }

        std::vector<uint8_t> response_pdu;
        response_pdu.push_back(slave_id_);
        response_pdu.push_back(0x02);

        uint8_t byte_count = static_cast<uint8_t>((quantity + 7) / 8);
        response_pdu.push_back(byte_count);

        for (uint16_t byte_index = 0; byte_index < byte_count; ++byte_index) {
            uint8_t discrete_byte = 0;
            for (uint8_t bit_index = 0; bit_index < 8; ++bit_index) {
                uint16_t discrete_index = start_address + byte_index * 8 + bit_index;
                if (discrete_index < start_address + quantity && discrete_index < data_.discrete_inputs.size()) {
                    if (data_.discrete_inputs[discrete_index]) {
                        discrete_byte |= (1 << bit_index);
                    }
                }
            }
            response_pdu.push_back(discrete_byte);
        }

        return addProtocolHeader(response_pdu, protocol, transaction_id);
    }

    std::vector<uint8_t> handleReadHoldingRegisters(const std::vector<uint8_t>& request, ProtocolType protocol, uint16_t transaction_id) {
        if (request.size() < 6) {
            return createExceptionResponse(0x03, 0x03, protocol, transaction_id);
        }

        uint16_t start_address = (request[2] << 8) | request[3];
        uint16_t quantity = (request[4] << 8) | request[5];

        if (quantity < 1 || quantity > 125) {
            return createExceptionResponse(0x03, 0x03, protocol, transaction_id);
        }

        if (!validateAddressRange(start_address, quantity, data_.holding_registers.size())) {
            return createExceptionResponse(0x03, 0x02, protocol, transaction_id);
        }

        std::vector<uint8_t> response_pdu;
        response_pdu.push_back(slave_id_);
        response_pdu.push_back(0x03);
        response_pdu.push_back(static_cast<uint8_t>(quantity * 2));

        for (uint16_t i = 0; i < quantity; ++i) {
            uint16_t register_value = data_.holding_registers[start_address + i];
            response_pdu.push_back(register_value >> 8);
            response_pdu.push_back(register_value & 0xFF);
        }

        return addProtocolHeader(response_pdu, protocol, transaction_id);
    }

    std::vector<uint8_t> handleReadInputRegisters(const std::vector<uint8_t>& request, ProtocolType protocol, uint16_t transaction_id) {
        if (request.size() < 6) {
            return createExceptionResponse(0x04, 0x03, protocol, transaction_id);
        }

        uint16_t start_address = (request[2] << 8) | request[3];
        uint16_t quantity = (request[4] << 8) | request[5];

        if (quantity < 1 || quantity > 125) {
            return createExceptionResponse(0x04, 0x03, protocol, transaction_id);
        }

        if (!validateAddressRange(start_address, quantity, data_.input_registers.size())) {
            return createExceptionResponse(0x04, 0x02, protocol, transaction_id);
        }

        std::vector<uint8_t> response_pdu;
        response_pdu.push_back(slave_id_);
        response_pdu.push_back(0x04);
        response_pdu.push_back(static_cast<uint8_t>(quantity * 2));

        for (uint16_t i = 0; i < quantity; ++i) {
            uint16_t register_value = data_.input_registers[start_address + i];
            response_pdu.push_back(register_value >> 8);
            response_pdu.push_back(register_value & 0xFF);
        }

        return addProtocolHeader(response_pdu, protocol, transaction_id);
    }

    std::vector<uint8_t> handleWriteSingleCoil(const std::vector<uint8_t>& request, ProtocolType protocol, uint16_t transaction_id) {
        if (request.size() < 6) {
            return createExceptionResponse(0x05, 0x03, protocol, transaction_id);
        }

        uint16_t address = (request[2] << 8) | request[3];
        uint16_t value = (request[4] << 8) | request[5];

        if (value != 0x0000 && value != 0xFF00) {
            return createExceptionResponse(0x05, 0x03, protocol, transaction_id);
        }

        if (address >= data_.coils.size()) {
            return createExceptionResponse(0x05, 0x02, protocol, transaction_id);
        }

        data_.coils[address] = (value == 0xFF00);

        // 返回相同的请求作为响应
        std::vector<uint8_t> response_pdu(request.begin() + 1, request.end());
        response_pdu[0] = slave_id_;
        return addProtocolHeader(response_pdu, protocol, transaction_id);
    }

    std::vector<uint8_t> handleWriteSingleRegister(const std::vector<uint8_t>& request, ProtocolType protocol, uint16_t transaction_id) {
        if (request.size() < 6) {
            return createExceptionResponse(0x06, 0x03, protocol, transaction_id);
        }

        uint16_t address = (request[2] << 8) | request[3];
        uint16_t value = (request[4] << 8) | request[5];

        if (address >= data_.holding_registers.size()) {
            return createExceptionResponse(0x06, 0x02, protocol, transaction_id);
        }

        data_.holding_registers[address] = value;

        std::vector<uint8_t> response_pdu(request.begin() + 1, request.end());
        response_pdu[0] = slave_id_;
        return addProtocolHeader(response_pdu, protocol, transaction_id);
    }

    std::vector<uint8_t> handleWriteMultipleCoils(const std::vector<uint8_t>& request, ProtocolType protocol, uint16_t transaction_id) {
        if (request.size() < 7) {
            return createExceptionResponse(0x0F, 0x03, protocol, transaction_id);
        }

        uint16_t start_address = (request[2] << 8) | request[3];
        uint16_t quantity = (request[4] << 8) | request[5];
        uint8_t byte_count = request[6];

        if (quantity < 1 || quantity > 0x07B0 || byte_count != (quantity + 7) / 8) {
            return createExceptionResponse(0x0F, 0x03, protocol, transaction_id);
        }

        if (!validateAddressRange(start_address, quantity, data_.coils.size())) {
            return createExceptionResponse(0x0F, 0x02, protocol, transaction_id);
        }

        // 写入线圈数据
        for (uint16_t i = 0; i < quantity; ++i) {
            uint8_t byte_index = static_cast<uint8_t>(i / 8);
            uint8_t bit_index = static_cast<uint8_t>(i % 8);
            if (byte_index < byte_count) {
                uint8_t coil_byte = request[7 + byte_index];
                data_.coils[start_address + i] = (coil_byte >> bit_index) & 0x01;
            }
        }

        // 构建响应
        std::vector<uint8_t> response_pdu;
        response_pdu.push_back(slave_id_);
        response_pdu.push_back(0x0F);
        response_pdu.push_back(request[2]); // 起始地址高字节
        response_pdu.push_back(request[3]); // 起始地址低字节
        response_pdu.push_back(request[4]); // 数量高字节
        response_pdu.push_back(request[5]); // 数量低字节

        return addProtocolHeader(response_pdu, protocol, transaction_id);
    }

    std::vector<uint8_t> handleWriteMultipleRegisters(const std::vector<uint8_t>& request, ProtocolType protocol, uint16_t transaction_id) {
        if (request.size() < 9) {
            return createExceptionResponse(0x10, 0x03, protocol, transaction_id);
        }

        uint16_t start_address = (request[2] << 8) | request[3];
        uint16_t quantity = (request[4] << 8) | request[5];
        uint8_t byte_count = request[6];

        if (quantity < 1 || quantity > 0x007B || byte_count != quantity * 2) {
            return createExceptionResponse(0x10, 0x03, protocol, transaction_id);
        }

        if (!validateAddressRange(start_address, quantity, data_.holding_registers.size())) {
            return createExceptionResponse(0x10, 0x02, protocol, transaction_id);
        }

        // 写入寄存器数据
        for (uint16_t i = 0; i < quantity; ++i) {
            uint16_t register_value = (request[7 + i * 2] << 8) | request[8 + i * 2];
            data_.holding_registers[start_address + i] = register_value;
        }

        // 构建响应
        std::vector<uint8_t> response_pdu;
        response_pdu.push_back(slave_id_);
        response_pdu.push_back(0x10);
        response_pdu.push_back(request[2]); // 起始地址高字节
        response_pdu.push_back(request[3]); // 起始地址低字节
        response_pdu.push_back(request[4]); // 数量高字节
        response_pdu.push_back(request[5]); // 数量低字节

        return addProtocolHeader(response_pdu, protocol, transaction_id);
    }

    std::vector<uint8_t> createExceptionResponse(uint8_t function_code, uint8_t exception_code, ProtocolType protocol, uint16_t transaction_id) {
        if (exception_callback_) {
            exception_callback_(slave_id_, function_code, exception_code);
        }

        std::vector<uint8_t> response_pdu;
        response_pdu.push_back(slave_id_);
        response_pdu.push_back(function_code | 0x80); // 设置异常标志
        response_pdu.push_back(exception_code);

        return addProtocolHeader(response_pdu, protocol, transaction_id);
    }

    std::vector<uint8_t> addProtocolHeader(const std::vector<uint8_t>& pdu, ProtocolType protocol, uint16_t transaction_id) {
        std::vector<uint8_t> frame;

        switch (protocol) {
        case ProtocolType::MODBUS_TCP: {
            // MBAP头: 事务标识符(2) + 协议标识符(2) + 长度(2) + 单元标识符(1)
            frame.push_back(transaction_id >> 8);
            frame.push_back(transaction_id & 0xFF);
            frame.push_back(0x00); // 协议标识符高字节
            frame.push_back(0x00); // 协议标识符低字节
            frame.push_back((pdu.size() >> 8) & 0xFF); // 长度高字节
            frame.push_back(pdu.size() & 0xFF); // 长度低字节
            frame.insert(frame.end(), pdu.begin(), pdu.end());
            break;
        }
        case ProtocolType::MODBUS_RTU: {
            frame = pdu;
            uint16_t crc = calculateCRC(frame);
            frame.push_back(crc & 0xFF);
            frame.push_back(crc >> 8);
            break;
        }
        case ProtocolType::MODBUS: {
            frame = pdu;
            break;
        }
        }

        return frame;
    }

    std::vector<uint8_t> removeProtocolHeader(const std::vector<uint8_t>& frame, ProtocolType protocol) {
        std::vector<uint8_t> pdu;

        switch (protocol) {
        case ProtocolType::MODBUS_TCP: {
            if (frame.size() >= 7) {
                pdu.assign(frame.begin() + 6, frame.end());
            }
            break;
        }
        case ProtocolType::MODBUS_RTU: {
            if (frame.size() >= 4) {
                // 移除CRC校验
                pdu.assign(frame.begin(), frame.end() - 2);
                // 验证CRC
                uint16_t received_crc = (frame[frame.size() - 1] << 8) | frame[frame.size() - 2];
                uint16_t calculated_crc = calculateCRC(pdu);
                if (received_crc != calculated_crc) {
                    return std::vector<uint8_t>(); // CRC校验失败
                }
            }
            break;
        }
        case ProtocolType::MODBUS: {
            pdu = frame;
            break;
        }
        }

        return pdu;
    }

    uint16_t calculateCRC(const std::vector<uint8_t>& data) {
        uint16_t crc = 0xFFFF;

        for (uint8_t byte : data) {
            crc ^= byte;
            for (int i = 0; i < 8; ++i) {
                if (crc & 0x0001) {
                    crc = (crc >> 1) ^ 0xA001;
                }
                else {
                    crc >>= 1;
                }
            }
        }

        return crc;
    }

    bool validateAddressRange(uint16_t address, uint16_t quantity, size_t max_size) {
        return (address < max_size) && (address + quantity <= max_size) && (quantity > 0);
    }

    private:
        uint8_t slave_id_;
        ModbusData& data_;
        ExceptionCallback exception_callback_;
};
